/*
 * CopyRight(C) Kenbin
 */

#include "lws_segment.h"
#include "lws_segger.h"

static LWS::CTrieDict *lws_triedict = NULL;
static LWS::CHmm *lws_person_hmm = NULL;
static LWS::CHmm *lws_location_hmm = NULL;
static LWS::CHmm *lws_postag_hmm = NULL;

static int lws_option_mask = 0;

char* lws_pos[] = {
    "Ag", "a", "ad", "an", "b", "c", "Dg", "d", "e", "f",
    "g", "h", "i", "j", "k", "l", "m", "Ng", "n", "nr",
    "ns", "nt", "nx", "nz", "o", "p", "q", "r", "s", "Tg", 
    "t", "u", "Vg", "v", "vd", "vn", "w", "x", "y", "z"
};

bool LWS::lws_init(const char *data_path, int option_mask)
{
    lws_option_mask = option_mask;

    lws_triedict = new CTrieDict;

    char trie_binfile[256] = {0};
    snprintf(trie_binfile, sizeof(trie_binfile), "%s/trie.dat", data_path);

    char CHAR_file[256] = {0};
    snprintf(CHAR_file, sizeof(CHAR_file), "%s/CHAR.txt", data_path);

    if(!lws_triedict->Load(trie_binfile, CHAR_file)) {
	return false;
    }

    if((option_mask & LWS_PERSON_RECOG) == LWS_PERSON_RECOG) {
	lws_person_hmm = new CHmm;

	char start_and_trans_file[256] = {0};
	snprintf(start_and_trans_file, sizeof(start_and_trans_file), 
		"%s/person_start_and_trans.txt", data_path);

	if(!lws_person_hmm->LoadStartAndTrans(start_and_trans_file, 10)) {
	    return false;
	}

	char emit_binfile[256] = {0};
	snprintf(emit_binfile, sizeof(emit_binfile), "%s/person_emit.dat", data_path);

	if(!lws_person_hmm->LoadEmit(emit_binfile)) {
	    return false;
	}
    }

    if((option_mask & LWS_LOCATION_RECOG) == LWS_LOCATION_RECOG) {
	lws_location_hmm = new CHmm;

	char start_and_trans_file[256] = {0};
	snprintf(start_and_trans_file, sizeof(start_and_trans_file), 
		"%s/location_start_and_trans.txt", data_path);

	if(!lws_location_hmm->LoadStartAndTrans(start_and_trans_file, 6)) {
	    return false;
	}

	char emit_binfile[256] = {0};
	snprintf(emit_binfile, sizeof(emit_binfile), "%s/location_emit.dat", data_path);

	if(!lws_location_hmm->LoadEmit(emit_binfile)) {
	    return false;
	}
    }

    if((option_mask & LWS_POSTAG) == LWS_POSTAG) {
	lws_postag_hmm = new CHmm;

	char start_and_trans_file[256] = {0};
	snprintf(start_and_trans_file, sizeof(start_and_trans_file), 
		"%s/postag_start_and_trans.txt", data_path);

	if(!lws_postag_hmm->LoadStartAndTrans(start_and_trans_file, 6)) {
	    return false;
	}

	char emit_binfile[256] = {0};
	snprintf(emit_binfile, sizeof(emit_binfile), "%s/postag_emit.dat", data_path);

	if(!lws_postag_hmm->LoadEmit(emit_binfile)) {
	    return false;
	}
    }

    return true;
}

void LWS::lws_uninit()
{
    delete lws_triedict;
    delete lws_person_hmm;
    delete lws_location_hmm;
    delete lws_postag_hmm;

    lws_triedict = NULL;
    lws_person_hmm = NULL;
    lws_location_hmm = NULL;
    lws_postag_hmm = NULL;

    lws_option_mask = 0;
}

void *LWS::lws_create_segger(int option)
{
    return new CSegment(lws_triedict, lws_person_hmm, lws_location_hmm, 
	    lws_postag_hmm, option & lws_option_mask);
}

void LWS::lws_destroy_segger(void *segger)
{
    delete (CSegment *)segger;
}

bool LWS::lws_segment(void *segger, const char *text)
{
    if(segger == NULL) {
	return false;
    }

    CSegment *seg = (CSegment *)segger;
    return seg->Segment(text);
}

int LWS::lws_word_count(void *segger)
{
    if(segger == NULL) {
	return 0;
    }

    CSegment *seg = (CSegment *)segger;
    return seg->WordCount();
}

char *LWS::lws_get_word(void *segger, int idx)
{
    if(segger == NULL) {
	return NULL;
    }

    CSegment *seg = (CSegment *)segger;
    if(idx >= seg->WordCount()) {
	return NULL;
    }

    return seg->WordPtrs()[idx];
}

char *LWS::lws_get_pos(void *segger, int idx)
{
    if(segger == NULL) {
	return NULL;
    }

    CSegment *seg = (CSegment *)segger;
    if(idx >= seg->WordCount()) {
	return NULL;
    }

    if(seg->WordSpecTags().find(idx) != seg->WordSpecTags().end()) {
	return lws_pos[36]; /* w */
    }

    return lws_pos[seg->BestPosTags()[idx]];
}

int LWS::lws_get_position(void *segger, int idx)
{
    if(segger == NULL) {
	return NULL;
    }

    CSegment *seg = (CSegment *)segger;
    if(idx >= seg->WordCount()) {
	return NULL;
    }

    return seg->WordPositions()[idx];
}
